package avicit.bdp.dds.api.utils.notebook;

import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @金航数码科技有限责任公司
 * @作者：developer
 * @邮箱：developer@avic-digital.com
 * @创建时间： 2022-05-23
 * @类说明：NoteBookTaskCacheManager
 * @修改记录：
 * @注意事项：
 * @主要功能：Jupyter NoteBook上下文缓存
 */
@Service
public class NoteBookTaskCacheManager {

    private static NoteBookTaskCacheManager instance = new NoteBookTaskCacheManager();

    private NoteBookTaskCacheManager() {

    }

    public static NoteBookTaskCacheManager getInstance() {
        return instance;
    }

    /**
     * 上下文缓存
     */
    private Map<String, NoteBookTaskContext> taskExecutionContextCache = new ConcurrentHashMap<>();

    /**
     * 根据NoteBook ID获取上下文
     * @param noteBookId
     * @return
     */
    public NoteBookTaskContext getByTaskInstanceId(String noteBookId) {
        if (StringUtils.isEmpty(noteBookId)) {
            return null;
        }

        return taskExecutionContextCache.get(noteBookId);
    }

    /**
     * 缓存NoteBook任务上下文
     * @param noteBookTaskContext
     */
    public void cacheTaskExecutionContext(NoteBookTaskContext noteBookTaskContext) {
        if (noteBookTaskContext == null) {
            return;
        }

        if (StringUtils.isNotBlank(noteBookTaskContext.getId())) {
            taskExecutionContextCache.put(noteBookTaskContext.getId(), noteBookTaskContext);
        }
    }

    /**
     * 移除NoteBook任务上下文
     * @param noteBookId
     */
    public void removeByTaskInstanceId(String noteBookId) {
        if (StringUtils.isEmpty(noteBookId)) {
            return;
        }

        if (taskExecutionContextCache.containsKey(noteBookId)) {
            taskExecutionContextCache.remove(noteBookId);
        }
    }
}
